'''
https://leetcode.cn/problems/subarrays-with-k-different-integers/
'''
from typing import List


class Solution:
    def subarraysWithKDistinct(self, nums: List[int], k: int) -> int:
        # 有多少个子数组，包含的元素个数小于等于K个
        n = len(nums)

        def num_of_le_K_different_numbers(k: int) -> int:
            cnt = [0] * (n + 1)
            l, r = 0, 0
            collect = 0  # 有多少个不同的数
            ans = 0
            while r < n:
                count = cnt[nums[r]] = cnt[nums[r]] + 1
                if count == 1:
                    collect += 1
                while l <= r and collect > k:
                    count = cnt[nums[l]] = cnt[nums[l]] - 1
                    l += 1
                    if count == 0:
                        collect -= 1
                ans += r - l + 1  # l...r内以r结尾的ans: l~r, l+1~r,l+2~r,...r~r
                r += 1
            return ans

        # 转移问题: 刚好为k个的 = num of (<= k个的) - num of (<= k-1 个的)
        return num_of_le_K_different_numbers(k) - num_of_le_K_different_numbers(k - 1)
